import gradio as gr
import numpy as np
from PIL import Image, ImageOps
from pinecone import Pinecone

# 初始化 Pinecone
pc = Pinecone("0b996cab-2029-4ed3-8646-20e9dec53ae2")

# 指定索引名称
index_name = 'mnist-index'
index = pc.Index(index_name)

# 定义预测函数
def predict_digit(image):
    # 如果 image 是字典，提取图像数据
    if isinstance(image, dict):
        image = image.get('composite')
        if image is None:
            raise ValueError("Image data not found in 'composite'")

    # 处理图像数据
    if image.shape[2] == 4:
        alpha_channel = image[:, :, 3]
    else:
        alpha_channel = np.mean(image[:, :, :3], axis=2)

    # 转换为 PIL 图像并调整大小
    pil_image = Image.fromarray(alpha_channel).resize((8, 8))
    pil_image = ImageOps.invert(pil_image.convert('L'))
    pil_image = pil_image.point(lambda x: 0 if x < 128 else 255, 'L')

    # 转换为 NumPy 数组并展平
    image_np = np.array(pil_image).reshape(1, -1)
    # 将像素值归一化到 [0, 1]
    image_np = image_np / 255.0

    # 在 Pinecone 中执行推理
    try:
        response = index.query(queries=image_np.tolist(), top_k=1)
        if response['matches']:
            prediction = response['matches'][0]['id']  # 获取预测结果
            return int(prediction)
        else:
            return "没有找到匹配结果"
    except Exception as e:
        return f"查询错误: {e}"

# 创建 Gradio 接口
with gr.Blocks() as demo:
    gr.Markdown("# 手写数字识别应用")

    with gr.Row():
        with gr.Column():
            sketchpad = gr.Sketchpad(label="Draw a digit", width=280, height=280)
            clear_button = gr.Button("Clear")
            submit_button = gr.Button("Submit")
        with gr.Column():
            label_output = gr.Label(label="Prediction Output")

    submit_button.click(fn=predict_digit, inputs=sketchpad, outputs=label_output)
    clear_button.click(lambda: None, None, sketchpad)

# 启动 Gradio 接口
demo.queue().launch()

